# =========== Copyright 2023 @ CAMEL-AI.org. All Rights Reserved. ===========
# ... (copyright header) ...
"""
这个文件定义了 ChatChain 类，它是整个 ChatAD 项目的核心编排和调度器。
它的主要职责是：
1.  加载所有必要的配置（如阶段、角色、任务）。
2.  管理整个对话流程的生命周期，从预处理到后处理。
3.  按顺序执行在配置文件中定义的每一个"阶段（Phase）"。
4.  初始化并管理一个共享的"聊天环境（ChatEnv）"，用于在不同阶段之间传递状态和数据。
5.  在任务结束后，收集并输出最终的产物。
"""

import importlib
import json
import logging
import os
import shutil
import time
from datetime import datetime

from camel.agents import RolePlaying
from camel.configs import ChatGPTConfig
from camel.typing import TaskType, ModelType
from chatad.chat_env import ChatEnv, ChatEnvConfig
from chatad.statistics import get_info
from camel.web_spider import modal_trans
from chatad.utils import log_visualize, now
from visualizer.app import log_file, send_msg
from chatad.data_classifier import DataClassifier

def check_bool(s):
    return s.lower() == "true"


class ChatChain:
    """
    ChatChain 类，用于管理和执行由多个AI智能体参与的多阶段对话任务。
    """
    def __init__(self,
                 config_path: str = None,
                 config_phase_path: str = None,
                 config_role_path: str = None,
                 task_prompt: str = None,
                 project_name: str = None,
                 org_name: str = None,
                 model_type: ModelType = ModelType.DEFAULT_MODEL) -> None:
        """
        初始化ChatChain实例。
        Args:
            config_path (str): ChatChain的通用配置文件路径。
            config_phase_path (str): 阶段（Phase）配置文件路径。
            config_role_path (str): 角色（Role）配置文件路径。
            task_prompt (str): 用户输入的原始任务描述。
            project_name (str): 本次任务的项目名称。
            org_name (str): 所属组织的名称。
            model_type (ModelType): 要使用的AI模型类型。
        """

        # 加载配置文件
        self.config_path = config_path
        self.config_phase_path = config_phase_path
        self.config_role_path = config_role_path
        self.project_name = project_name
        self.org_name = org_name
        self.model_type = model_type

        with open(self.config_path, 'r', encoding="utf8") as file:
            self.config = json.load(file)  # ChatChainConfig.json
        with open(self.config_phase_path, 'r', encoding="utf8") as file:
            self.config_phase = json.load(file) # PhaseConfig
        with open(self.config_role_path, 'r', encoding="utf8") as file:
            self.config_role = json.load(file)  # RoleConfig

        # 初始化ChatChain配置和招募
        self.chain = self.config["chain"]
        self.recruitments = self.config["recruitments"]
        self.web_spider = self.config["web_spider"]

        # 初始化默认最大聊天轮数
        self.chat_turn_limit_default = 10

        # 初始化ChatEnv
        self.chat_env_config = ChatEnvConfig(clear_structure=check_bool(self.config["clear_structure"]),
                                             background_prompt=self.config["background_prompt"])
                                             
        self.chat_env = ChatEnv(self.chat_env_config)

        # 用户输入提示将进行自我改进（如果设置 "self_improve": "True" 在 ChatChainConfig.json 中）
        # 自我改进在 self.preprocess 中完成
        self.task_prompt_raw = task_prompt
        self.task_prompt = ""

        # 初始化角色提示
        self.role_prompts = dict()
        for role in self.config_role:
            self.role_prompts[role] = "\n".join(self.config_role[role])

        # 初始化日志
        self.start_time, _ = self.get_logfilepath()
        self.log_filepath=log_file

        # 初始化SimplePhase实例
        # 从 chatad.phase 导入 PhaseConfig.json 中定义的所有 SimplePhases
        # 注意：在 PhaseConfig.json 中只存在 SimplePhases，ComposedPhases 在 ChatChainConfig.json 中定义并将在 self.execute_step 中导入
        # ComposedPhases are defined in ChatChainConfig.json and will be imported in self.execute_step
        self.phase_module = importlib.import_module("chatad.phase")
        self.phases = dict()
        for phase in self.config_phase:
            assistant_role_name = self.config_phase[phase]['assistant_role_name']
            user_role_name = self.config_phase[phase]['user_role_name']
            phase_prompt = "\n\n".join(self.config_phase[phase]['phase_prompt'])
            phase_class = getattr(self.phase_module, phase)
            phase_instance = phase_class(assistant_role_name=assistant_role_name,
                                         user_role_name=user_role_name,
                                         phase_prompt=phase_prompt,
                                         role_prompts=self.role_prompts,
                                         phase_name=phase,
                                         model_type=self.model_type,
                                         log_filepath=self.log_filepath)
            self.phases[phase] = phase_instance

    def make_recruitment(self):
        """
        "人员招募"阶段：根据配置，实例化所有需要的AI智能体角色。
        """
        for employee in self.recruitments:
            self.chat_env.recruit(agent_name=employee)

    def execute_step(self, phase_item: dict):
        """
        执行链中的单个阶段
        Args:
            phase_item: ChatChainConfig.json 中的单个阶段配置
        Returns:
        """

        phase = phase_item['phase']
        phase_type = phase_item['phaseType']
        # 对于 SimplePhase，只需从 self.phases 中查找并执行 "Phase.execute" 方法
        if phase_type == "SimplePhase":
            max_turn_step = phase_item['max_turn_step']
            need_reflect = check_bool(phase_item['need_reflect'])
            if phase in self.phases:
                self.chat_env = self.phases[phase].execute(self.chat_env,
                                                           self.chat_turn_limit_default if max_turn_step <= 0 else max_turn_step,
                                                           need_reflect)
            else:
                raise RuntimeError(f"Phase '{phase}' is not yet implemented in chatad.phase")
        else:
            raise RuntimeError(f"PhaseType '{phase_type}' is not yet implemented.")

    def execute_chain(self):
        """
        【核心】执行对话链。
        该方法会按顺序迭代执行配置文件中定义的所有阶段（Phase）。
        """
        for phase_item in self.chain:
            self.execute_step(phase_item)

    def get_logfilepath(self):
        """
        获取日志路径（在软件路径下）
        返回值
            start_time：开始制作软件的时间
            log_filepath：日志路径

        """
        start_time = now()
        filepath = os.path.dirname(__file__)
        # root = "/".join(filepath.split("/")[:-1])
        root = os.path.dirname(filepath)
        # directory = root + "/WareHouse/"
        directory = os.path.join(root, "logs")   # WareHouse
        log_filepath = os.path.join(directory,
                                    "{}.log".format("_".join([self.project_name, self.org_name, start_time])))
        return start_time, log_filepath

    def pre_processing(self):
        """
        执行任务开始前的预处理工作。
        """
        filepath = os.path.dirname(__file__)
        root = os.path.dirname(filepath)
        directory = os.path.join(root, "WareHouse")

        if self.chat_env.config.clear_structure:  # 去除WareHouse文件夹下非项目内且非log文件
            for filename in os.listdir(directory):
                file_path = os.path.join(directory, filename)
                if os.path.isfile(file_path) and not filename.endswith(".log"):
                    os.remove(file_path)
                    print("{} Removed.".format(file_path))

        ware_path = os.path.join(directory, "_".join([self.project_name, self.org_name, self.start_time]))
        self.chat_env.set_directory(ware_path)  # /WareHouse下的路径  ware_path：项目文件夹名


        # 将配置文件复制到软件路径
        shutil.copy(self.config_path, ware_path)
        shutil.copy(self.config_phase_path, ware_path)
        shutil.copy(self.config_role_path, ware_path)

        # 将任务提示写入软件
        with open(os.path.join(ware_path, self.project_name + ".prompt"), "w") as f:
            f.write(self.task_prompt_raw)

        # 任务初始化
        preprocess_msg = "**[Preprocessing]**\n\n"
        chat_gpt_config = ChatGPTConfig()
        
        # 记录配置信息
        preprocess_msg += "**ChatAD Starts** ({})\n\n".format(self.start_time)
        preprocess_msg += "**Timestamp**: {}\n\n".format(self.start_time)
        preprocess_msg += "**config_path**: {}\n\n".format(self.config_path)
        preprocess_msg += "**config_phase_path**: {}\n\n".format(self.config_phase_path)
        preprocess_msg += "**config_role_path**: {}\n\n".format(self.config_role_path)
        preprocess_msg += "**task_prompt**: {}\n\n".format(self.task_prompt_raw)
        preprocess_msg += "**project_name**: {}\n\n".format(self.project_name)
        preprocess_msg += "**Log File**: {}\n\n".format(self.log_filepath)
        preprocess_msg += "**ChatADConfig**:\n{}\n\n".format(self.chat_env.config.__str__())
        preprocess_msg += "**ChatGPTConfig**:\n{}\n\n".format(chat_gpt_config)
        log_visualize(preprocess_msg)

        # 初始化任务提示
        if check_bool(self.config['self_improve']):  # 任务提升
            self.chat_env.env_dict['task_prompt'] = self.self_task_improve(self.task_prompt_raw)
        else:
            self.chat_env.env_dict['task_prompt'] = self.task_prompt_raw
        if(check_bool(self.web_spider)):  # 联网搜索
            self.chat_env.env_dict['task_description'] = modal_trans(self.task_prompt_raw)

        # 仅获取可用数据类别，不加载数据
        data_classifier = DataClassifier()
        available_categories = data_classifier.get_available_categories()
        self.chat_env.env_dict["available_categories"] = available_categories
        
        log_visualize(f"**[数据准备]** 已获取 {len(available_categories)} 个可用数据类别，接下来将进行分类选择")
        
        # 发送到前端显示 - 仅显示可用类别列表
        send_msg("系统", f"数据系统已准备完毕，可用的数据类别有 {len(available_categories)} 个，包括：\n\n" + 
                 "\n".join(available_categories))

    def post_processing(self):
        """
        执行任务结束后的后处理工作。
        """
        self.chat_env.write_meta()
        filepath = os.path.dirname(__file__)
        root = os.path.dirname(filepath)

        post_info = "**[Post Info]**\n\n"
        # 计算任务执行时间
        now_time = now()
        time_format = "%Y%m%d%H%M%S%f"
        datetime1 = datetime.strptime(self.start_time, time_format)
        datetime2 = datetime.strptime(now_time, time_format)
        duration = (datetime2 - datetime1).total_seconds()

        post_info += "Software Info: {}".format(
            get_info(self.chat_env.env_dict['directory'], self.log_filepath) + "\n\n🕑**duration**={:.2f}s\n\n".format(
                duration))

        post_info += "ChatAD Starts ({})".format(self.start_time) + "\n\n"
        post_info += "ChatAD Ends ({})".format(now_time) + "\n\n"

        directory = self.chat_env.env_dict['directory']
        if self.chat_env.config.clear_structure:
            for filename in os.listdir(directory):
                file_path = os.path.join(directory, filename)
                if os.path.isdir(file_path) and file_path.endswith("__pycache__"):  # 可能要去掉
                    shutil.rmtree(file_path, ignore_errors=True)
                    post_info += "{} Removed.".format(file_path) + "\n\n"

        log_visualize(post_info)

        logging.shutdown()
        time.sleep(1)

        shutil.move(self.log_filepath,
                    os.path.join(root, "WareHouse", "_".join([self.project_name, self.org_name, self.start_time]),
                                 os.path.basename(self.log_filepath)))
#可开启
    def self_task_improve(self, task_prompt):
        """
        询问代理改进用户查询提示
        Args:
            task_prompt: 原始用户查询提示

        Returns:
            revised_task_prompt: 提示工程师代理修改后的提示

        """
        self_task_improve_prompt = """我将给你一个简短的广告设计要求描述，请把它改写成一个详细的提示，使大型语言模型知道如何根据这个提示把这个广告做得更好,
提示应确保 LLM 制作的广告能够正确运行，这是您需要考虑的最重要的部分。请记住，修改后的提示不应超过 200 字，以下是简要说明:\"{}\". 
你修改后的内容应该以"<好的>"开头的消息，不要返回其他格式的消息。""".format(
            task_prompt)
        role_play_session = RolePlaying(
            assistant_role_name="PromptEngineer",
            assistant_role_prompt="你是一个专业的提示工程师，可以改进用户输入的提示，使大型语言模型更好地理解这些提示。",
            user_role_prompt="你是一个想要使用大型语言模型构建广告的用户。",
            user_role_name="User",
            task_type=TaskType.CHATAD,
            task_prompt="Do prompt engineering on user query",
            # with_task_specify=False,
            model_type=self.model_type,
        )
        _, input_user_msg = role_play_session.init_chat(None, None, self_task_improve_prompt)
        assistant_response, user_response = role_play_session.step(input_user_msg, True) # user_response为空
        # 此处assistant only设为true是因为对话单向，只需要PromptEngineer这个assistant回答就行
        revised_task_prompt = assistant_response.msg.content.split("<好的>")[-1].lower().strip()
        log_visualize(role_play_session.assistant_agent.role_name, "**PromptEngineer<->User on task_improve: , turn 0**"+"[ ]\n\n"+assistant_response.msg.content)
        log_visualize(
            "**[Task Prompt Self Improvement]**\n**原始用户任务**: {}\n**改进后任务**: {}".format(
                task_prompt, revised_task_prompt))
        return revised_task_prompt